#! /usr/bin/env python
'''
GUI - RotaPy

Interactive GUI for rotating seismograms and analysing
sensor orientation based on envelope functions.

Maximum envelope ratio (radial/transverse) should occur at a backazimuth
that is correct. Orientation correct for sensors is then calculated as the
difference between the starting backazimuth and the calculated backazimuth.
'''
import os.path
import sys
import string
import fnmatch

import numpy as np

from PyQt4.QtCore import *
from PyQt4.QtGui import *
from obspy.core import read
from matplotlib.figure import Figure
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt4agg import NavigationToolbar2QTAgg as NavigationToolbar
from matplotlib.widgets import MultiCursor
from obspy.signal import rotate
from obspy.signal import filter

__version__ = "0.3"

# Define some global values here
MIN_ANG = 0 # Minimum backazimuth angle for slider
MAX_ANG = 359 # Maximum backazimuth angle for slider
THRESHOLD = 50 # Threshold for zero'ing spikes
MUTE_WINDOW = 5 # No. of points to zero if spike is detected
REAL_TIME = True # Set False if on a slow-machine (determines if real-time update is allowed)

class MainWindow(QMainWindow):
    def __init__(self, parent=None):
        """
        Called when initializing the application
        DO NOT TOUCH!
        """
        super(MainWindow, self).__init__(parent)
        
        # Initialize UI
        self.initMainWindow()
        self.initPlot()
        self.initDock()
        self.initMenu()
        self.initSignal()
        self.center()
        
    def center(self):
        """
        Function that finds the center of the screen
        and places the main window there (don't touch this!)
        """
        screen = QDesktopWidget().screenGeometry()
        size = self.geometry()
        self.move( (screen.width() - size.width())/2, (screen.height() - size.height())/2 )
            
    def initMainWindow(self):
        # Initialize main window widget
        self.setWindowTitle("RotaPy") # Because it sounds better than PyRoti! 
        self.resize(600, 400)
        self.setMinimumSize(600, 400) # Set minimum window sizes
        self.inputdir = os.getcwd() # Set 'starting' directory as the current directory
        
        # Create dummy widget and setup grid layout for all items
        self.mainWidget = QWidget(self)
        self.setCentralWidget(self.mainWidget)
        self.mainLayout = QGridLayout(self.mainWidget)
        
        # Setup status bar displaying information about the loaded traces (at bottom)
        self.status = self.statusBar()
        self.seriesTitle = QLabel("Trace length [s]: ")
        self.seriesData = QLabel("--")
        self.samprateTitle = QLabel("Sampling Rate [Hz]")
        self.samprateData = QLabel("--")
        self.seriesData.setFrameStyle(QFrame.StyledPanel|QFrame.Sunken)
        self.samprateData.setFrameStyle(QFrame.StyledPanel|QFrame.Sunken)        
        
        self.status.addPermanentWidget(self.seriesTitle)
        self.status.addPermanentWidget(self.seriesData)
        self.status.addPermanentWidget(self.samprateTitle)
        self.status.addPermanentWidget(self.samprateData)
        self.status.showMessage("Ready", 1000)
    
    def initPlot(self):
        """
        Initiate the barebones plot and starting titles/labels.
        Only modify tick positions/labels etc
        """
        # Setup plotting window and slider
        self.bazNew = self.bazOrig = 1 # Set correct/modified back azimuth to '1'
        
        # Inititate plot figure (2x3 plots) and make them share x-axis (and y-axis)
        self.fig = Figure()
        self.ax1 = self.fig.add_subplot(231)
        self.ax2 = self.fig.add_subplot(232, sharex=self.ax1)
        self.ax3 = self.fig.add_subplot(233, sharex=self.ax1)
        self.ax4 = self.fig.add_subplot(234, sharex=self.ax1, sharey=self.ax1)
        self.ax5 = self.fig.add_subplot(235,sharex=self.ax1)
        self.ax6 = self.fig.add_subplot(236,sharex=self.ax1)
        
        self.fig.subplots_adjust(hspace=0.1,left=0.1,right=0.95,wspace=0) # Adjust the whitespace surrounding plots
        
        # Set misc plot params (titles, tick positions)
        self.ax1.set_ylabel("BAZ: ") # Original data plot
        self.ax4.set_ylabel("BAZ: ") # Modified data plot
        self.ax4.set_xlabel("Time [s]")
        self.ax1.set_title("No data loaded") # Nothing loaded
        
        self.ax1.tick_params(labelbottom=False,labelleft=False)
        self.ax4.tick_params(labelleft=False,labelbottom=False)
        self.ax3.tick_params(labelleft=False,labelbottom=False)
        self.ax2.tick_params(labelleft=False,labelbottom=False)
        self.ax5.tick_params(labelleft=False,labelbottom=False)
        self.ax6.tick_params(labelleft=False,labelbottom=False)
        
        self.canvas = FigureCanvas(self.fig) # Create the displayed figure canvas

        self.mainLayout.addWidget(self.canvas,0,0,10,11) # Add canvas to main window
        self.navToolBar = NavigationToolbar(self.canvas, self.mainWidget) # Create canvas toolbar
        
        # Create azimuth slider 
        self.slider = QSlider(Qt.Horizontal)
        self.slider.setRange(MIN_ANG, MAX_ANG) # Set range between 0, 359
        self.slider.setTickInterval(30) # Tickmark every 30 degrees
        self.slider.setTickPosition(1) # Ticks above the slider bar
        self.slider.setValue(self.bazOrig) # Set starting value to '1'
        self.slider.setTracking(REAL_TIME) # Checks global value to see if real-time update is turned on

        # Box for displaying angle next to slider
        self.angBox = QLabel()
        self.angBox.setText(str(self.slider.value()))
        
        # Setup check boxes for component display
        self.checkBoxZ = QCheckBox("Vertical")
        self.checkBoxR = QCheckBox("Radial")
        self.checkBoxT = QCheckBox("Transverse")
        self.checkBoxEnv = QCheckBox("Envelope")
        # Set initial states to 'checked'
        self.checkBoxZ.setChecked(True)
        self.checkBoxR.setChecked(True)
        self.checkBoxT.setChecked(True)
        self.checkBoxEnv.setChecked(True)
        # Add widgets to main window - Work on adding these to
        # sub-layouts with expanding/shrinking to make best use of space
        self.mainLayout.addWidget(self.slider,10,4,1,6)
        self.mainLayout.addWidget(self.checkBoxZ,10,0,1,1)
        self.mainLayout.addWidget(self.checkBoxR,10,1,1,1)
        self.mainLayout.addWidget(self.checkBoxT,10,2,1,1)
        self.mainLayout.addWidget(self.checkBoxEnv,10,3,1,1)
        self.mainLayout.addWidget(self.angBox)
        
        self.canvas.draw() # Update plot window

    def initDock(self):
        """
        Setup the file explorer (don't touch this!)
        """
        # Setup dock widgets  (file explorer)
        self.fileListWidget = QListWidget()
        self.fileListWidget.setMaximumWidth(150) # Set minimum width to 150px
        
        self.fileBrowserWidget = QDockWidget("File Browser", self)
        self.fileBrowserWidget.setObjectName("FileBrowserWidget")
        self.fileBrowserWidget.setAllowedAreas(Qt.LeftDockWidgetArea)
        self.fileBrowserWidget.setFeatures(QDockWidget.NoDockWidgetFeatures)
        self.fileBrowserWidget.setWidget(self.fileListWidget)
        # Add the dock to the lefthand side
        self.addDockWidget(Qt.LeftDockWidgetArea, self.fileBrowserWidget)
        
    def initMenu(self):
        """
        Sets up the menubar and adds it (don't touch this!)
        """
        # Setup menubar
        filePrefAction = self.createAction("&Preferences", self.inputPref,
                "Ctrl+Z", "prefnew", "Input preferences")
        fileQuitAction = self.createAction("&Quit", self.close,
                "Ctrl+Q", "filequit", "Close the application")
        fileScanAction = self.createAction("&Scan...", self.scanDir,
                "Ctrl+S", "filescan", "Scan for files")
        filePrintAction = self.createAction("&Print", self.openPrintDialog,
                "Ctrl+P", "fileprint", "Print the plot window")
        # Setup plot menu
        zoomPlotAction = self.createAction("&Zoom..", self.navToolBar.zoom,
                "z", "plotzoom", "Zoom data")
        panPlotAction = self.createAction("&Pan...", self.navToolBar.pan,
                "p", "plotpan", "Pan data")
        resetPlotAction = self.createAction("&Home..", self.navToolBar.home,
                "h", "plotreset", "Reset figure")
        cursorPlotAction = self.createAction("Cursor", self.multiCursor,
                "c", "plotcursor", "Switch to multi-cursor")
        
        self.fileMenu = self.menuBar().addMenu("&File")
        self.plotMenu = self.menuBar().addMenu("&Plot")
        
        self.addActions(self.fileMenu,(fileScanAction,filePrefAction,None,filePrintAction,None,fileQuitAction))
        self.addActions(self.plotMenu,(cursorPlotAction,zoomPlotAction,panPlotAction,
                                       resetPlotAction))
    
    def initSignal(self):
        """
        Links interactions with the GUI to the function calls (don't touch this!)
        """
        # Setup slots/signals (that turn clicks/keyboard entries to functions!)
        self.connect(self.slider, SIGNAL("valueChanged(int)"),self.updateData) # If slider changes, update display
        self.connect(self.fileListWidget, SIGNAL("itemDoubleClicked(QListWidgetItem*)"), self.loadData) # If file is double clicked, load data
        self.connect(self.checkBoxZ,SIGNAL("stateChanged(int)"), self.updateData) # if checkbox changes, update data
        self.connect(self.checkBoxR,SIGNAL("stateChanged(int)"), self.updateData) # if checkbox changes, update data
        self.connect(self.checkBoxT,SIGNAL("stateChanged(int)"), self.updateData) # if checkbox changes, update data
        self.connect(self.checkBoxEnv,SIGNAL("stateChanged(int)"), self.updateData) # if checkbox changes, update data
    
    def updateData(self):
        """
         This part of the code does all the calculations and is called by most of the 
         GUI signals (checkboxes, etc)
        """
        self.bazOld = self.bazNew # Store the modified BAZ as 'old' so we can check if its the same.
        self.bazNew = self.slider.value() # Get modified azimuth value from from slider  
        self.angBox.setText(str(self.bazNew)) # Update the label to show the modified azimuth
        
        # If the modified BAZ is same as last value or if this is fresh data, perform the calculations
        # otherwise just refresh the plotted lines.
        if (self.bazNew != self.bazOld) or self.clean:
            # Update the modified BAZ label on the plot
            self.ylabel2 = "BAZ: %i" % (self.bazNew) 
            self.ax4.set_ylabel(self.ylabel2)
            #####################################
            ### Perform all calculations here ###
            #####################################
            # Calculate modified radial, transverse components
            self.radNew, self.transNew = rotate.rotate_NE_RT(self.seisN, self.seisE, self.bazNew)
            # Calculate vertical envelope
            self.envVert = filter.envelope(self.seisZ)
            # Calculate radial envelope
            self.envRadNew = filter.envelope(self.radNew)
            # Calculate transverse envelope
            self.envTransNew = filter.envelope(self.transNew)
            # Calculate energy of vertical envelope
            self.envVertEnergy = np.square(np.abs(self.envVert))
            # Calculate energy of radial envelope
            self.envRadEnergy = np.square(np.abs(self.envRadNew))
            # Calculate energy of transverse envelope
            self.envTransEnergy = np.square(np.abs(self.envTransNew))
            # Calcaulate energy of radial component
            self.radEnergy = np.square(np.abs(self.radNew))
            # Calcaulate energy of transverse component
            self.transEnergy = np.square(np.abs(self.transNew))
            # Calcaulate energy of vertical component
            self.vertEnergy = np.square(np.abs(self.seisZ))
            # Calculate ratios of radial/transverse components and their envelopes
            self.ratio = np.true_divide(self.radEnergy,self.transEnergy)
            self.envRatio = np.true_divide(self.envRadEnergy,self.envTransEnergy)
            
            # Mute first 4.0 seconds (since all data begins around 4.5 - 5 s)
            # then zero any data points above the threshold value + window
            # This is a horribly crude spike suppression method    
            self.ratio[self.t<4.0] = 0
            self.envRatio[self.t<4.0] = 0
        
            for i in np.arange(len(self.envRatio)):
                if self.envRatio[i] > THRESHOLD:
                    self.envRatio[i:i+MUTE_WINDOW] = 0
                if self.ratio[i] > THRESHOLD:
                    self.ratio[i:i+MUTE_WINDOW] = 0
    
        # Either way, let's do the following:        
        # Erase all plotted lines and refresh data
        del self.ax1.lines[:], self.ax2.lines[:], self.ax3.lines[:],\
         self.ax4.lines[:], self.ax5.lines[:], self.ax6.lines[:]
  
        alpha = 0.8 # Transparency value
        
        # All plotting checks are performed below
        # i.e. which boxes are ticked, what to display?
        
        # If Vertical checkbox is ticked (2)
        if self.checkBoxZ.checkState() == 2:
            self.ax1.plot(self.t,self.seisZ,'k',alpha=alpha)
            self.ax3.plot(self.t,self.vertEnergy,'k',alpha=alpha)
            self.ax4.plot(self.t,self.seisZ,'k',alpha=alpha)
        # If Radial checkbox is ticked (2)    
        if self.checkBoxR.checkState() == 2:
            self.ax1.plot(self.t,self.rad,'r',alpha=alpha)
            self.ax4.plot(self.t,self.radNew,'r',alpha=alpha)
            self.ax3.plot(self.t,self.radEnergy,'r',alpha=alpha)
        # If Transverse checkbox is ticked (2)    
        if self.checkBoxT.checkState() == 2:    
            self.ax1.plot(self.t,self.trans,'b',alpha=alpha)
            self.ax3.plot(self.t,self.transEnergy,'b',alpha=alpha)
            self.ax4.plot(self.t,self.transNew,'b',alpha=alpha)
        # If Envelope checkbox is ticked and vertical is ticked (2)
        if self.checkBoxEnv.checkState() == 2 and self.checkBoxZ.checkState() == 2:
            self.ax4.plot(self.t,self.envVert,'k:',alpha=alpha)
            self.ax2.plot(self.t,self.envVertEnergy,'k:',alpha=alpha)
        # If Envelope checkbox is ticked and radial is ticked (2)   
        if self.checkBoxEnv.checkState() == 2 and self.checkBoxR.checkState() == 2:
            self.ax4.plot(self.t,self.envRadNew,'r:',alpha=alpha)
            self.ax2.plot(self.t,self.envRadEnergy,'r:',alpha=alpha)
        # If Envelope checkbox is ticked and transverse is ticked (2)      
        if self.checkBoxEnv.checkState() == 2 and self.checkBoxT.checkState() == 2:
            self.ax4.plot(self.t,self.envTransNew,'b:',alpha=alpha)
            self.ax2.plot(self.t,self.envTransEnergy,'b:',alpha=alpha)
        # If radial and transverse are both ticked, plot the ratios (this may be uncool if you want to examine one at a time)
        if self.checkBoxR.checkState() == 2 and self.checkBoxT.checkState() == 2:
            self.ax5.plot(self.t, self.envRatio, 'k:',alpha=alpha)
            self.ax6.plot(self.t,self.ratio, 'k',alpha=alpha)
        
        # Re-draw the plot
        self.canvas.draw()

    def loadData(self):
        """
        Function to work out which file we clicked on in the explorer
        then load the Z/N/E traces using obspy read() function
        
        Then stack all the data traces for each component into seisZ/seisN/seisE variables
        and label plot accordingly, then call updateData() to calculate the various items of interst
        """
        # If the filelist has at least one item clicked, continue
        if self.fileListWidget.currentRow() != -1:
            # Ask us for the backazimuth ands set both modified/original value to this
            self.bazNew = self.bazOrig = self.openBazDialog()
            self.clean = True # Set a flag to show this data is unmodified
            
            #Obtain the file by matching where you click to the file names
            file = self.fileInfo[self.fileListWidget.currentRow()][1]
            # Update status bar to show we're loading the data
            self.status.showMessage("Loading data...")
            
            # Try to load all data and setup the initial plot parameters
            try:
                self.stZ = read(file.replace('*','Z')) # Load vertical
                self.stN = read(file.replace('*','N')) # Load north
                self.stE = read(file.replace('*','E')) # Load east
            
                # Go through the loaded data and stack all traces
                # If this is the first file, then get the delta/npts and update the statusbar
                for i, tr in enumerate(self.stZ):
                    if i == 0:
                        self.delta = tr.stats.delta
                        self.npts = tr.stats.npts
                        self.seriesData.setText(str(self.npts*self.delta))
                        self.samprateData.setText(str(1./self.delta))
                        self.t = np.arange(0, self.delta*(self.npts-0.5), self.delta)
                        self.seisZ = tr.data
                    else:   
                        self.seisZ += tr.data
                # Do it again for the 'north' data    
                for i, tr in enumerate(self.stN):
                    if i == 0:
                        self.seisN = tr.data
                    else:
                        self.seisN += tr.data
                # Do it again for the 'east' data
                for i, tr in enumerate(self.stE):
                    if i == 0:
                        self.seisE = tr.data
                    else:
                        self.seisE += tr.data
                # Calculate the unmodified 'correct' radial/transverse components (only need to do this once)
                self.rad, self.trans = rotate.rotate_NE_RT(self.seisN, self.seisE, self.bazOrig)
                # Show data has finished loading
                self.status.showMessage("Loading complete!", 3000)
                
                # Update plot with station, line, and backazimuth information
                info = os.path.basename(file).split('.')
                plttitle = "%s - %s" % (info[0], info[1])
                
                # Reset the axes when loading new data
                self.ax1.cla()
                self.ax2.cla()
                self.ax3.cla()
                self.ax4.cla()
                self.ax5.cla()
                self.ax6.cla()
                
                self.ylabel1 = "BAZ: %i" % (self.bazOrig)
                self.ylabel2 = "BAZ: %i" % (self.bazNew)
                self.ax1.set_ylabel(self.ylabel1)
                self.ax4.set_ylabel(self.ylabel2)
                self.ax2.set_xlabel("Energy (envelope)")
                self.ax5.set_xlabel("Ratio (R/T) (envelope)")
                self.ax3.set_xlabel("Energy")
                self.ax6.set_xlabel("Ratio (R/T)")
                self.ax1.set_title(plttitle)
                
                # Call updateData() via slider-value change (very hacky!)
                # I.e. changing the value will call the valueChanged() signal, which calls updateData()
                self.slider.setValue(self.bazNew) # Set the new value
                self.angBox.setText(str(self.slider.value())) # Update the label to show new value
                
            # If loading data failed, display an unhelpful error message and do nothing  
            except Exception:
                self.status.showMessage("ERROR: Loading failed!", 5000)
                
    def scanDir(self):
        """
        Scan the directory we choose and all sub-directories (recursive)
        for appropriate SEG-Y files. If a single component is missing, none 
        of the matching SEG-Y files will be displayed at the end
        """
        # Get input directory
        self.inputdir = self.inputPref() # Bring up the directory choice window
        self.extension = '*.sgy' # Could change this for SEG-Y, mSEED, SAC
        
        # Find all files and create a list of tuples showing:
        # (display name, filename)
        self.fileInfo = [('.'.join(os.path.basename(f).split('.')[0:4]), f) \
                         for f in self.Walk(self.inputdir,pattern=self.extension)]
        
        # Re-fresh the file explorer by removing and re-adding all entries
        for item in range(self.fileListWidget.count()):
            self.fileListWidget.takeItem(item)
        self.fileListWidget.addItems([f[0] for f in self.fileInfo])
  
    def Walk(self, root, recurse=1, pattern='*.sgy', return_folders=0,skipdir=''):
        """
        From ActiveState.com Recipe: 52664-flexible-directory-walking/
        Scan from a root directory recursively for all files matching our pattern
        """
        # initialize
        result = []
        # must have at least root folder
        try:
            names = os.listdir(root)
            for dir in skipdir.split(','):
                if dir in names:
                    names.remove(dir)
        except os.error:
            return result
        # expand pattern
        pattern = pattern or '*'
        pat_list = string.splitfields( pattern , ',' )
        # check each file
        for name in names:
            fullname = os.path.normpath(os.path.join(root, name))
            # grab if it matches our pattern and entry type
            for pat in pat_list:
                if fnmatch.fnmatch(name, pat):
                    if os.path.isfile(fullname) or (return_folders and os.path.isdir(fullname)):
                        result.append(fullname)
                    continue
            # recursively scan other folders, appending results
            if recurse:
                if os.path.isdir(fullname) and not os.path.islink(fullname):
                    result = result + self.Walk( fullname, recurse, pattern, return_folders )      
        # Delete duplicates (convert to set, then back to list)
        return sorted(list(set([self.wildcardFile(f) for f in result])))
    
    def wildcardFile(self, file):
        """
        Replace the channel identifier in a filename with '*'
        this helps us create the displayname/filename master list
        """
        f = os.path.split(file)
        tmp = f[1].split('.')
        tmp[-2] = '*'
        
        return os.path.join(f[0],'.'.join(tmp))  
   
    def createAction(self, text, slot=None, shortcut=None, icon=None,
                     tip=None, checkable=False, signal="triggered()"):
        """
        Convenience method for setting up QActions (don't touch this!)
        """
        action = QAction(text, self)
        if icon is not None:
            action.setIcon(QIcon(":/{0}.png".format(icon)))
        if shortcut is not None:
            action.setShortcut(shortcut)
        if tip is not None:
            action.setToolTip(tip)
            action.setStatusTip(tip)
        if slot is not None:
            self.connect(action, SIGNAL(signal), slot)
        if checkable:
            action.setCheckable(True)
        return action

    def addActions(self, target, actions):
        """
        Convenience method for adding QActions to a menu (don't touch this!)
        """
        for action in actions:
            if action is None:
                target.addSeparator()
            else:
                target.addAction(action)
                
    def inputPref(self):
        """
        Brings up the window to let you choose which directory you want to scan for files
        """
        if self.inputdir == os.getcwd():
            return str(QFileDialog.getExistingDirectory(self, "Select Data Directory"))
        else:
            return str(QFileDialog.getExistingDirectory(self, "Select Data Directory", self.inputdir))
    
    def openBazDialog(self):
        """
        Open up the window and ask for back-azimuth (correct) value.
        Accepts between 0-359 only
        """
        dialog = QDialog(self)
        
        dialog.setWindowTitle("Backazimuth")
        dialog.setSizeGripEnabled(False)
        
        okButton = QPushButton("&Ok")
        bazLabel = QLabel("Backazimuth: ")
        bazEdit = QLineEdit()
        bazRange = QIntValidator(MIN_ANG,MAX_ANG,bazEdit)
        bazEdit.setValidator(bazRange)
        bazEdit.setText(str(MIN_ANG))
        bazEdit.selectAll()
        
        dialog.connect(okButton, SIGNAL("clicked()"),dialog.close)

        grid = QGridLayout()
        grid.addWidget(bazLabel,0,0,1,1)
        grid.addWidget(bazEdit,0,1,1,1)
        grid.addWidget(okButton,1,0,2,1)
        
        dialog.setLayout(grid)
        dialog.exec_()
        # Return the inputted backazimuth as integer
        return int(bazEdit.text())
    
    def openPrintDialog(self):
        dialog = QPrintDialog(self)
        dialog.exec_()
    
    def multiCursor(self):
        """
        Adds the vertical cursor that links position between all plots
        Useful for comparing across figures (don't touch this)
        """
        if hasattr(self, 'multi'):
            del self.multi
            self.canvas.draw()
        else:
            self.multi = MultiCursor(self.canvas,(self.ax1, self.ax2, self.ax3,self.ax4,self.ax5, self.ax6),
                                      color='r', lw=1)
            self.canvas.draw()
 
if __name__ == "__main__":
    # Create and run the application!
    app = QApplication(sys.argv)
    form = MainWindow()
    form.show()
    app.exec_()